import sys
import inspect
import typing
import io
from typing import Any, Callable, Dict, Optional, Type, TypeVar, List

from pydantic import BaseModel, Field, create_model

from calc_manager.base.utils import load_module_from_path
from calc_manager.base.fields import FieldInfo, OPTIONAL_KEY, REQUIRED_KEY
from calc_manager.base.plugins import PluginInvokeParams
from calc_manager.plugins.plugins_manager import PluginManager
from calc_manager.plugins.plugin_object import PluginHelper

import logging
logger = logging.getLogger(__name__)


class OperatorInfo(BaseModel):
    operator_id: str
    operator_name: str
    operator_path: str
    operator_workingdir: Optional[str] = None
    operator_func_entrypoint: str
    operator_desc: Optional[str] = None
    #make class hashable
    class Config:
        frozen = True


# TODO: add un_load_model_function
class Operator:
    PARAMTER_CLASS = "Parameter"
    RESULT_CLASS = "Result"
    INVOKE_PARAM_KEY = "param"
    INVOKE_PLUGINS_KEY = "plugins"
    INVOKE_RESULT_KEY = "result"

    def __init__(self, operator_info: OperatorInfo):
        self.info = operator_info
        self.m = self.load_model()

    def dispose(self):
        """
        Dispose of the operator by unloading all loaded modules from memory.
        """
        if hasattr(self, 'm'):
            # Get all attributes of the loaded module
            module_attrs = dir(self.m)
            
            # Remove all references to objects defined in the module
            for attr in module_attrs:
                if not attr.startswith('__'):
                    delattr(self.m, attr)
            
            # Remove the module from sys.modules if it's there
            if self.m.__name__ in sys.modules:
                del sys.modules[self.m.__name__]
            
            # Delete the reference to the module in the operator
            del self.m
        
        # Force garbage collection to clean up any remaining references
        import gc
        gc.collect()
        
        logger.info(f"Disposed operator: {self.info.operator_name}")

    def load_model(self):
        logger.info("Loading operator model from path: %s, working_dir %s, operator_name: %s",
                    self.info.operator_path, self.info.operator_workingdir, self.info.operator_name)
        try:
            if self.info.operator_workingdir:
                working_dir = (
                    self.info.operator_workingdir
                    if self.info.operator_workingdir[-1] == "/"
                    else self.info.operator_workingdir + "/"
                )
                m = load_module_from_path(
                    working_dir + self.info.operator_path, self.info.operator_name
                )
            else:
                m = load_module_from_path(self.info.operator_path, self.info.operator_name)
            return m
        except Exception as e:
            raise Exception("Failed to load operator model: %s" % e)

    def get_attr_in_model(self, attr_name):
        return getattr(self.m, attr_name)

    def extract_params(self) -> dict[str, FieldInfo]:
        return extract_model_info(
            self.get_attr_in_model(Operator.PARAMTER_CLASS)
        )  ##TODO: refactor

    def extract_result(self) -> dict[str, FieldInfo]:
        return extract_model_info(
            self.get_attr_in_model(Operator.RESULT_CLASS)
        )  ##TODO: refactor

    def get_entrypoint_function(self):
        return getattr(self.m, self.info.operator_func_entrypoint)

    @classmethod
    def get_plugin_objects_from_argv(cls, plugin_manager: PluginManager, plugins_invoke_argv: Dict[str, Dict[str, Any]]) -> Dict[str, PluginHelper]:
        a = plugin_manager.get_plugin_params_invoke_params_by_param_dict(plugins_invoke_argv)
        return Operator._set_plugin_param_with_extra_argv_values(plugin_manager, plugins_invoke_argv, a)

    def get_plugin_objects_from_cls_name_obj_dict(self, plugin_cls_name_to_obj: Dict[str, Any]) -> Dict[str, Any]:
        current_operator_invoke_argv_sig = inspect.signature(self.get_entrypoint_function())
        return {k: v for k, v in plugin_cls_name_to_obj.items() if k in current_operator_invoke_argv_sig.parameters}

    def get_operator_argv_plugin_object(self, plugin_manager: PluginManager, plugins_invoke_argv: Dict[str, Dict[str, Any]]) -> Dict[str, PluginHelper]:
        # this version filter the plugin params by current operator's invoke params
        # not all plugins will be loaded
        current_operator_invoke_argv_sig = inspect.signature(self.get_entrypoint_function())
        a = plugin_manager.get_plugin_params_invoke_params_by_param_dict(plugins_invoke_argv)
        a = {k: v for k, v in a.items() if v.plugin_cls_param_name in current_operator_invoke_argv_sig.parameters}
        return Operator._set_plugin_param_with_extra_argv_values(plugin_manager, plugins_invoke_argv, a)
    
    def get_operator_plugin_invoke_params(self, plugin_manager: PluginManager, exclude_plugin_cls_param_name: List[str] = []) -> Dict[str, List[FieldInfo]]:
        """
        Get plugin helper objects for all plugin parameters required by the operator's entrypoint function.

        This method inspects the operator's entrypoint function signature and creates plugin helper objects
        for each plugin parameter. It validates that the plugin parameter names are registered with the
        plugin manager.

        Args:
            plugin_manager (PluginManager): The plugin manager instance to handle plugin operations

        Returns:
            Dict[str, PluginInvokeParams]: Dictionary mapping plugin names (not plugin class parameter names) 
                                          to their helper objects. The keys are the actual plugin names 
                                          (e.g. "my_plugin") rather than the class parameter names 
                                          (e.g. "plugin1")

        Raises:
            ValueError: If a plugin parameter name in the function signature is not registered with 
                       the plugin manager
        """
        ret = {}
        # ignore the invoke param key
        exclude_plugin_cls_param_name.append(Operator.INVOKE_PARAM_KEY)
        current_operator_invoke_argv_sig = inspect.signature(self.get_entrypoint_function())
        for plugin_cls_param_name, _v in current_operator_invoke_argv_sig.parameters.items():
            if plugin_cls_param_name in exclude_plugin_cls_param_name:
                continue
            if not plugin_manager.plugin_cls_param_name_is_valid(plugin_cls_param_name):
                raise ValueError(f"Plugin class param name {plugin_cls_param_name} is not valid in plugin manager")
            ret[plugin_manager.get_plugin_name_by_cls_param_name(plugin_cls_param_name)] = plugin_manager.get_plugin_fields_by_plugin_cls_param_name(plugin_cls_param_name)
        return ret
    
    @classmethod
    def get_valid_param_invoke(cls, func, values: Dict[str, Any]) -> Dict[str, Any]:
        """
        Pick up the valid parameters from values and check the type of the parameters
        """
        sig = inspect.signature(func)
        ret = {}
        for param_name, param in sig.parameters.items():
            if param_name not in values:
                raise ValueError(f"Parameter '{param_name}' is required but was not provided in values.")
            
            expected_type = param.annotation
            provided_value = values[param_name]
            provided_type = type(provided_value)
            
            if not isinstance(provided_value, expected_type):
                raise TypeError(f"Parameter '{param_name}' is expected to be of type '{expected_type.__name__}', but got '{provided_type.__name__}'.")
            else:
                ret[param_name] = provided_value
        return ret


    @classmethod
    def get_param_class_and_create_instance(
        cls, func, param_name: str, values: Dict[str, Any]
    ):
        sig = inspect.signature(func)
        #if param_name in sig.parameters:
        param_type = sig.parameters[param_name].annotation

        if issubclass(param_type, BaseModel):
            return param_type(**values)
        else:
            raise TypeError("Parameter type is not a subclass of BaseModel")

    #TODO: move to plugin manager not here
    @classmethod
    def _set_plugin_param_with_extra_argv_values(cls, plugin_manager: PluginManager, invoke_param_argv: Dict[str, Dict[str, Any]], plugins_invoke_param: Dict[str, PluginInvokeParams]) -> Dict[str, PluginHelper]:
        """
        Set plugin parameters with extra argument values.

        This method processes plugin parameters and creates plugin helper objects based on the provided arguments.

        Args:
            plugin_manager (PluginManager): The plugin manager instance to handle plugin operations
            invoke_param_argv (Dict[str, Dict[str, Any]]): Dictionary containing original invoke parameters
            plugins_invoke_param (Dict[str, PluginInvokeParams]): Dictionary mapping plugin names to their invoke parameters

        Returns:
            Dict[str, PluginHelper]: Dictionary mapping plugin class parameter names to their helper objects, key is plugin class param name

        Raises:
            ValueError: If plugin key type is not str, plugin value type is not PluginInvokeParams,
                       or if plugin parameter name already exists in invoke_param_argv
        """
        invoke_param_orig = invoke_param_argv
        ret: Dict[str, PluginHelper] = {}
        logger.debug("Start to set plugin param with extra_argv values: %s", plugins_invoke_param)
        for plugin_name, plugin_invoke_param in plugins_invoke_param.items():
            if not isinstance(plugin_name, str):
                raise ValueError(f"Invalid plugin key type: expected str, got {type(plugin_name)}")
            if not isinstance(plugin_invoke_param, PluginInvokeParams):
                raise ValueError(f"Invalid plugin value type for {plugin_name}: expected PluginInvokeParams, got {type(plugin_invoke_param)}")
            
            plugin_cls_param_name = plugin_invoke_param.plugin_cls_param_name
            p = plugin_invoke_param.param
            plugin_object_helper = plugin_manager.get_plugin_object_helper_by_cls_param_name(plugin_cls_param_name, **p)
            # if plugin_cls_param_name not in invoke_param_orig:
            #     raise ValueError(f"Plugin param name {plugin_cls_param_name} already exists in invoke_raw data {invoke_param_orig}, Register plugin with name and invoke with plugin-name not plugin-param-name")
            ret[plugin_cls_param_name] = plugin_object_helper
        logger.debug("Set plugin param with extra_argv values: %s, from orig param_input %s", ret, invoke_param_orig)
        return ret

    # if only argc is used, then the function will be called with the parameters
    def invoke_func(self, *argc, **argv):
        """
        Invokes the entry point function of the operator with the provided parameters.

        This method prioritizes the use of `argc` if it is present and the first element is a dictionary.
        If `argc` is not provided or is not a dictionary, it falls back to using the dictionary provided
        under the `param_key` in `argv`.

        Note: if argv is used, argc is ignroe

        Parameters:
        *argc: tuple
            Positional arguments, where the first element can be a dictionary containing parameter values.
        **argv: dict
            Keyword arguments, where `param_key` should map to a dictionary of parameter values.

        Raises:
        ValueError
            If neither `argc` nor `argv` contains a valid dictionary of parameters.

        Returns:
        Any
            The result of invoking the entry point function with the constructed parameter object.

        Usage:
        - If a dictionary is provided as the first element in `argc`, it will be used as the parameter dictionary.
        - If no dictionary is found in `argc`, the method looks for a dictionary under `param_key` in `argv`.
        - If neither condition is met, a ValueError is raised..
        """

        param_key = Operator.INVOKE_PARAM_KEY
        func = self.get_entrypoint_function()
        logger.debug("Start to invoke operator: %s, function signature: %s" , self.info, inspect.signature(func))

        if param_key in argv and isinstance(argv[param_key], dict):
            param_obj = Operator.get_param_class_and_create_instance(
                func, param_key, argv[param_key]
            )
        elif argc and isinstance(argc[0], dict) and param_key not in argv:
            param_obj = Operator.get_param_class_and_create_instance(
                func, param_key, argc[0]
            )
        else:
            raise ValueError("Invalid parameter type. No valid param found in argv['%s'] or argc[0]. Expected a dictionary, but received argv: %s, argc: %s" % (param_key, argv, argc))

        argv.pop(param_key, None)
        argv[param_key] = param_obj
        #invoke_param = {param_key: param_obj}

        logger.info("Invoke function with param keys and types: %s", [(k, type(v).__name__) for k, v in argv.items()])

        argv = Operator.get_valid_param_invoke(func, argv)
        
        res = func(**argv)
        if hasattr(res, '__dict__'):
            attrs = vars(res)
            logger.info("Operator invoke function result - Number of attributes: %d, Attribute names: %s", len(attrs), list(attrs.keys()))
        else:
            logger.info("Operator invoke function result: %s", res)
        return res


def extract_model_info(model) -> Dict[str, FieldInfo]:
    """
    Extracts field information from a Pydantic model.

    This function analyzes a given Pydantic model and extracts detailed information about its fields,
    including their names, types, optionality, default values, and descriptions.

    Parameters:
    model (Type[BaseModel]): The Pydantic model to extract information from.

    Returns:
    Dict[str, FieldInfo]: A dictionary where keys are field names and values are FieldInfo objects
                          containing detailed information about each field.

    The function performs the following steps:
    1. Extracts field annotations and model fields from the input model.
    2. Defines a helper function `get_type_name` to convert complex types (like Union, List, etc.) to string representations.
    3. Iterates through each field in the model, determining:
       - Whether the field is optional
       - The field's type (converted to a string representation)
       - The field's description (if any)
       - The field's default value (if any)
    4. Constructs a FieldInfo object for each field with the gathered information.

    Note:
    - This function handles complex types, including nested types and Optional fields.
    - It uses global constants OPTIONAL_KEY and REQUIRED_KEY to denote field optionality.
    """
    field_info = {}
    annotations = model.__annotations__
    model_fields = model.model_fields

    def get_type_name(typ) -> str:
        origin = typing.get_origin(typ)
        if origin:
            if origin is typing.Union:
                args = ", ".join(
                    get_type_name(arg)
                    for arg in typing.get_args(typ)
                    if arg is not type(None)
                )
                return "Optional[%s]" % args
            else:
                args = ", ".join(get_type_name(arg) for arg in typing.get_args(typ))
                return "%s[%s]" % (origin.__name__, args)
        else:
            return typ.__name__

    for name, typ in annotations.items():
        is_optional = typing.get_origin(typ) == typing.Union and type(
            None
        ) in typing.get_args(typ)
        field_type = typing.get_args(typ)[0] if is_optional else typ
        field_type_str = get_type_name(field_type)
        description = (
            model_fields[name].description if model_fields[name].description else None
        )
        field_info[name] = FieldInfo(
            name=name,
            type=field_type_str,
            optional_field=OPTIONAL_KEY if is_optional else REQUIRED_KEY,
            description=description,
            default_val=(
                model_fields[name].default
                if model_fields[name].default is not Ellipsis
                else None
            ),
        )
    return field_info

## not use here just remove
#def create_param_class_from_func(
#    func, attr_name: str, fields_dict: Dict[str, FieldInfo]
#) -> Type[BaseModel]:
#    sig = inspect.signature(func)
#    param = sig.parameters.get(attr_name)
#    if not param:
#        raise ValueError(f"Parameter {attr_name} not found in function signature.")
#
#    field_definitions = {}
#    for field_name, field_info in fields_dict.items():
#        field_type = eval(field_info.type)
#        if (
#            field_info.optional_field == OPTIONAL_KEY
#        ):  ## TODO: refactor need imporve this logic
#            field_definitions[field_name] = (
#                Optional[field_type],
#                Field(default=None, description=field_info.description),
#            )
#        else:
#            field_definitions[field_name] = (
#                field_type,
#                Field(description=field_info.description),
#            )
#
#    DynamicParameter = create_model(attr_name, __base__=BaseModel, **field_definitions)
#    return create_model("DynamicParameter", **{attr_name: (DynamicParameter, ...)})


def print_operator_info(operator: Operator) -> str:
    output = io.StringIO()

    output.write("----BaseInfo-------\n")
    output.write(f"Operator Name: {operator.info.operator_name}\n")
    output.write(f"Operator Path: {operator.info.operator_path}\n")
    output.write(
        f"Operator Function Entrypoint: {operator.info.operator_func_entrypoint}\n"
    )
    output.write(
        f"Operator Description: {operator.info.operator_desc if operator.info.operator_desc else 'No Description'}\n"
    )

    output.write("----Parameters-----\n")
    params_info = operator.extract_params()
    if params_info:
        for field_name, field in params_info.items():
            output.write(
                f"  {field_name}: Type={field.type}, Optional={field.optional_field}, Default={field.default_val}, Description={field.description}\n"
            )
    else:
        output.write("  No parameters specified.\n")

    output.write("----Result---------\n")
    result_info = operator.extract_result()
    if result_info:
        for field_name, field in result_info.items():
            output.write(
                f"  {field_name}: Type={field.type}, Optional={field.optional_field}, Default={field.default_val}, Description={field.description}\n"
            )
    else:
        output.write("  No result fields specified.\n")

    output.write("-------------------\n")

    return output.getvalue()
